package com.lndsnx.ai.service.impl;

import com.lndsnx.ai.service.FileService;
import com.lndsnx.entity.Constant;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.reader.JsonReader;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.reader.jsoup.JsoupDocumentReader;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Service;

import java.util.List;

/**
 * @author: lndsnx
 * @date: 2025/7/8 21:10
 * @Version: 1.0
 * @description:
 */
@Slf4j
@Service
public class FileServiceImpl implements FileService {

    @Autowired
    private ChatClient chatClient;
    @Autowired
    private VectorStore vectorStore;

    @Override
    public boolean save(String chatId, Resource resource) {
        DocumentReader reader = getDocument(resource);
        List<Document> documents = reader.read();
        vectorStore.add(documents);
        return true;
    }

    @Override
    public String chat(String query, String chatId) {
        return chatClient.prompt(query)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId))
                .call().content();
    }

    @Override
    public String chatRagAdvisor(String query, String chatId) {
        RetrievalAugmentationAdvisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .vectorStore(vectorStore)
                        .build())
                .build();

        return chatClient.prompt(query)
                .advisors(retrievalAugmentationAdvisor)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, chatId))
                .call().content();
    }

    // TODO 待完善 数据少没问题，数据多就出错了，待优化
    private DocumentReader getDocument(Resource resource) {
        String filename = resource.getFilename();
        String fileType = filename.substring(filename.lastIndexOf("."));
        DocumentReader reader;
        if(Constant.JSON_TYPE.equalsIgnoreCase(fileType)){
            reader = new JsonReader(resource);
        }else if(Constant.TEXT_TYPE.equalsIgnoreCase(fileType)){
            reader = new TextReader(resource);
        }else if(Constant.MARKDOWN_TYPE.equalsIgnoreCase(fileType)){
            // TODO 有病哦，跟其他长不一样，先使用 TikaDocumentReader
            // reader = new MarkdownDocumentReader(resource);
            reader = new TikaDocumentReader(resource);
        }else if(Constant.PDF_TYPE.equalsIgnoreCase(fileType)){
            reader = new PagePdfDocumentReader(resource);
        }else if(Constant.HTML_TYPE.equalsIgnoreCase(fileType)){
            reader = new JsoupDocumentReader(resource);
        }else {
            reader = new TikaDocumentReader(resource);
        }
        return reader;
    }
}
